import cv2
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt
from pathlib import Path
import time


class VehicleDetector:
    """车辆识别检测器类，用于检测图像和视频中的车辆"""

    def __init__(self, model_path='yolov8n.pt', confidence_threshold=0.5):
        """
        初始化车辆检测器

        参数:
            model_path: YOLO模型路径，默认使用预训练的yolov8n模型
            confidence_threshold: 置信度阈值，过滤低置信度检测结果
        """
        # 加载YOLO模型
        self.model = YOLO(model_path)

        # 车辆相关类别（COCO数据集）
        self.vehicle_classes = {
            2: 'car',  # 汽车
            3: 'motorcycle',  # 摩托车
            5: 'bus',  # 公交车
            7: 'truck'  # 卡车
        }

        # 置信度阈值
        self.confidence_threshold = confidence_threshold

        # 为不同车辆类型定义不同颜色
        self.colors = {
            'car': (0, 255, 0),  # 绿色
            'motorcycle': (0, 0, 255),  # 红色
            'bus': (255, 0, 0),  # 蓝色
            'truck': (255, 255, 0)  # 青色
        }

    def detect_vehicles(self, image):
        """
        检测图像中的车辆

        参数:
            image: 输入图像（BGR格式）

        返回:
            检测结果列表，每个元素包含边界框坐标、置信度和车辆类型
        """
        # 运行模型预测
        results = self.model(image, conf=self.confidence_threshold)

        # 提取检测结果
        detections = []
        for result in results:
            for box in result.boxes:
                # 获取类别ID
                class_id = int(box.cls)

                # 只保留车辆相关类别
                if class_id in self.vehicle_classes:
                    # 获取边界框坐标
                    x1, y1, x2, y2 = map(int, box.xyxy[0])

                    # 获取置信度
                    confidence = float(box.conf)

                    # 获取车辆类型
                    vehicle_type = self.vehicle_classes[class_id]

                    detections.append({
                        'bbox': (x1, y1, x2, y2),
                        'confidence': confidence,
                        'type': vehicle_type
                    })

        return detections

    def draw_detections(self, image, detections):
        """
        在图像上绘制检测结果

        参数:
            image: 输入图像
            detections: 检测结果列表

        返回:
            绘制了检测结果的图像
        """
        # 创建图像副本，避免修改原图
        image_copy = image.copy()

        for detection in detections:
            x1, y1, x2, y2 = detection['bbox']
            vehicle_type = detection['type']
            confidence = detection['confidence']

            # 绘制边界框
            cv2.rectangle(image_copy, (x1, y1), (x2, y2), self.colors[vehicle_type], 2)

            # 绘制标签
            label = f"{vehicle_type}: {confidence:.2f}"
            cv2.putText(image_copy, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors[vehicle_type], 2)

        return image_copy

    def process_image(self, image_path, output_path=None):
        """
        处理单张图像并保存结果

        参数:
            image_path: 输入图像路径
            output_path: 输出图像路径，若为None则不保存

        返回:
            处理后的图像
        """
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")

        # 检测车辆
        start_time = time.time()
        detections = self.detect_vehicles(image)
        end_time = time.time()

        print(f"检测耗时: {end_time - start_time:.4f}秒")
        print(f"检测到{len(detections)}辆车辆")

        # 绘制检测结果
        result_image = self.draw_detections(image, detections)

        # 保存结果
        if output_path:
            # 创建输出目录（如果不存在）
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            cv2.imwrite(output_path, result_image)
            print(f"结果已保存至: {output_path}")

        return result_image

    def process_video(self, video_path, output_path=None, display=True):
        """
        处理视频并保存结果

        参数:
            video_path: 输入视频路径
            output_path: 输出视频路径，若为None则不保存
            display: 是否实时显示处理结果
        """
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"无法打开视频: {video_path}")

        # 获取视频属性
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # 初始化视频编写器（如果需要保存输出）
        out = None
        if output_path:
            # 创建输出目录（如果不存在）
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        frame_count = 0
        start_time = time.time()

        try:
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break

                frame_count += 1

                # 检测车辆并绘制结果
                detections = self.detect_vehicles(frame)
                result_frame = self.draw_detections(frame, detections)

                # 显示处理进度
                if frame_count % 10 == 0:
                    progress = (frame_count / total_frames) * 100
                    print(f"处理进度: {progress:.1f}% ({frame_count}/{total_frames})")

                # 写入输出视频
                if out:
                    out.write(result_frame)

                # 显示结果
                if display:
                    cv2.imshow('Vehicle Detection', result_frame)

                    # 按'q'键退出
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break

        finally:
            # 释放资源
            cap.release()
            if out:
                out.release()
            cv2.destroyAllWindows()

        # 计算处理时间
        end_time = time.time()
        processing_time = end_time - start_time
        print(f"视频处理完成，耗时: {processing_time:.2f}秒")
        print(f"帧率: {frame_count / processing_time:.2f} FPS")

        if output_path:
            print(f"结果已保存至: {output_path}")


if __name__ == "__main__":
    # 创建车辆检测器实例
    detector = VehicleDetector(confidence_threshold=0.4)

    # 示例：处理单张图像
    try:
        # 替换为你的图像路径
        image_result = detector.process_image(
            image_path="test_vehicle.jpg",
            output_path="output/detected_vehicle.jpg"
        )

        # 显示结果（使用matplotlib，支持Jupyter环境）
        plt.figure(figsize=(10, 6))
        plt.imshow(cv2.cvtColor(image_result, cv2.COLOR_BGR2RGB))
        plt.axis('off')
        plt.title('Vehicle Detection Result')
        plt.show()
    except Exception as e:
        print(f"处理图像时出错: {e}")

    # 示例：处理视频（取消注释以使用）
    """
    try:
        detector.process_video(
            video_path="test_video.mp4",
            output_path="output/detected_video.mp4",
            display=True
        )
    except Exception as e:
        print(f"处理视频时出错: {e}")
    """
